"""
API客户端 - 负责处理与外部API的交互
"""
import requests
import json
import uuid
import time
import urllib3
from typing import List, Dict, Any, Optional, Tuple
from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry

# 禁用不安全连接警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

def create_retry_session(retries=3, backoff_factor=0.3, 
                        status_forcelist=(500, 502, 504), 
                        session=None):
    """
    创建带有重试机制的会话
    
    参数:
        retries: 重试次数
        backoff_factor: 退避因子
        status_forcelist: 需要重试的HTTP状态码列表
        session: 现有会话对象
        
    返回:
        带有重试机制的会话对象
    """
    session = session or requests.Session()
    retry = Retry(
        total=retries,
        read=retries,
        connect=retries,
        backoff_factor=backoff_factor,
        status_forcelist=status_forcelist,
    )
    adapter = HTTPAdapter(max_retries=retry)
    session.mount('http://', adapter)
    session.mount('https://', adapter)
    return session

def generate_request_id(prefix: str = "") -> str:
    """
    生成唯一的请求ID
    
    参数:
        prefix: 请求ID的前缀
    
    返回:
        唯一的请求ID字符串
    """
    if prefix:
        return f"{prefix}-{str(uuid.uuid4())[:8]}"
    return str(uuid.uuid4())

def search_knowledge_chunks(question: str, top_k: int = 50) -> List[str]:
    """
    调用检索API获取与问题相关的知识块
    
    参数:
        question: 问题文本
        top_k: 返回的知识块数量
        
    返回:
        知识块列表
    """
    # API接口地址
    url = "https://iop.midea.com/rag-flow/api/search/keyword"
    
    # 请求头
    headers = {
        'appId': 'iop-knowledge',
        'appChannel': 'test',
        'Content-Type': 'application/json'
    }
    
    # 准备请求体
    payload = {
        "requestId": generate_request_id(question[:8]),
        "query": question,
        "languageList": ["zh", "en"],
        "userName": "chenmy2",
        "topk": top_k,
        "rerankMethod": "none",
        "skillList": [
            {
                "recallTypeList": ["text", "vector"],
                "skillId": "SKILL_CS31SK20250423000000192220",
                "topk": 10,
                "sourceSystem": "knowledge",
                "threshold": 0.0,
                "conditionItem": None,
                "faqAnswerConditionItem": None,
                "catalogIdList": []
            },
            {
                "recallTypeList": ["text", "vector"],
                "skillId": "SKILL_CS31SK20250418000000124477",
                "topk": 10,
                "sourceSystem": "ihr_selfhelp",
                "threshold": 0.0,
                "conditionItem": None,
                "faqAnswerConditionItem": None,
                "catalogIdList": []
            },
            {
                "recallTypeList": ["text", "vector"],
                "skillId": "SKILL_CS30SK20250417000000176056",
                "topk": 10,
                "sourceSystem": "official_doc",
                "threshold": 0.0,
                "conditionItem": None,
                "faqAnswerConditionItem": None,
                "catalogIdList": []
            },
            {
                "recallTypeList": ["text", "vector"],
                "skillId": "SKILL_CS30SK20250417000000182249",
                "topk": 10,
                "sourceSystem": "g-qms",
                "threshold": 0.0,
                "conditionItem": None,
                "faqAnswerConditionItem": None,
                "catalogIdList": []
            },
            {
                "recallTypeList": ["text", "vector"],
                "skillId": "1889572155196022785",
                "topk": 10,
                "sourceSystem": "knowledge_book",
                "threshold": 0.0,
                "conditionItem": None,
                "faqAnswerConditionItem": None,
                "catalogIdList": []
            },
            {
                "recallTypeList": ["text", "vector"],
                "skillId": "1879815395340005377",
                "topk": 10,
                "sourceSystem": "knowledge_book",
                "threshold": 0.0,
                "conditionItem": None,
                "faqAnswerConditionItem": None,
                "catalogIdList": []
            }
        ],
        "knowTypeList": None,
        "unAuthFilter": True,
        "faqSameKnowNeedOne": True,
        "lexiconGeneralizationList": None
    }
    
    try:
        # 创建带有重试机制的会话
        session = create_retry_session(retries=3)
        
        # 发送请求
        print(f"为问题 '{question}' 检索 {top_k} 个知识块")
        response = session.post(url, headers=headers, json=payload, verify=False, timeout=30)
        response.raise_for_status()  # 如果请求失败，抛出异常
        
        # 解析响应
        result = response.json()
        
        # 检查响应状态
        if result["code"] != "0":
            print(f"API调用失败: {result['msg']}")
            return []
        
        # 提取知识块内容（text字段）
        chunks = []
        for item in result["data"]:
            if "text" in item and len(item["text"]) > 0:
                chunks.append(item["text"][0])
        
        print(f"成功检索到 {len(chunks)} 个知识块")
        return chunks
        
    except requests.exceptions.SSLError as e:
        print(f"SSL连接错误: {str(e)}")
        print("尝试禁用SSL验证...")
        try:
            # 禁用SSL验证再次尝试
            response = requests.post(url, headers=headers, json=payload, verify=False, timeout=30)
            response.raise_for_status()
            result = response.json()
            
            if result["code"] != "0":
                print(f"API调用失败: {result['msg']}")
                return []
            
            chunks = []
            for item in result["data"]:
                if "text" in item and len(item["text"]) > 0:
                    chunks.append(item["text"][0])
            
            print(f"成功检索到 {len(chunks)} 个知识块")
            return chunks
        except Exception as e2:
            print(f"第二次尝试失败: {str(e2)}")
            return []
    except Exception as e:
        print(f"API调用异常: {str(e)}")
        return []

def evaluate_chunk_with_llm(question: str, chunk: str) -> Tuple[float, str]:
    """
    使用大模型对单个知识块进行评分
    
    参数:
        question: 问题文本
        chunk: 知识块内容
        
    返回:
        元组 (评分, 评分理由)
    """
    # API接口地址和请求头
    url = "https://aigc.midea.com/dify/server/v1/workflows/run"
    headers = {
        'Authorization': 'Bearer app-aPQkRmCQw81D4yCL4BcldMnF',
        'Content-Type': 'application/json'
    }
    
    # 准备请求体
    payload = {
        "inputs": {
            "question": question,
            "knowledge_chunk": chunk
        },
        "response_mode": "blocking",
        "user": "chenqh114"
    }
    
    try:
        # 创建带有重试机制的会话
        session = create_retry_session(retries=2)
        
        # 发送请求
        response = session.post(url, headers=headers, json=payload, verify=False, timeout=30)
        response.raise_for_status()
        
        # 解析响应
        result = response.json()
        
        if result.get("data", {}).get("status") == "succeeded":
            # 提取评分和理由
            output_str = result.get("data", {}).get("outputs", {}).get("result", "{}")
            output_json = json.loads(output_str)
            
            score = float(output_json.get("score", 0.0))
            reason = output_json.get("reason", "无评分理由")
            
            return score, reason
        else:
            error_msg = result.get("data", {}).get("error", "未知错误")
            print(f"API调用失败: {error_msg}")
            return 0.0, f"API调用失败: {error_msg}"
            
    except requests.exceptions.SSLError as e:
        print(f"评分SSL错误: {str(e)}")
        try:
            # 禁用SSL验证再次尝试
            response = requests.post(url, headers=headers, json=payload, verify=False, timeout=30)
            response.raise_for_status()
            result = response.json()
            
            if result.get("data", {}).get("status") == "succeeded":
                output_str = result.get("data", {}).get("outputs", {}).get("result", "{}")
                output_json = json.loads(output_str)
                
                score = float(output_json.get("score", 0.0))
                reason = output_json.get("reason", "无评分理由")
                
                return score, reason
            else:
                error_msg = result.get("data", {}).get("error", "未知错误")
                return 0.0, f"API调用失败: {error_msg}"
        except Exception as e2:
            return 0.0, f"评分重试失败: {str(e2)}"
    except Exception as e:
        print(f"评分过程异常: {str(e)}")
        return 0.0, f"评分异常: {str(e)}"

def evaluate_chunks_batch(question: str, chunks: List[str]) -> List[Tuple[str, float, str]]:
    """
    批量对知识块进行评分
    
    参数:
        question: 问题文本
        chunks: 知识块列表
        
    返回:
        元组列表，每个元组包含 (知识块, 评分, 理由)
    """
    print(f"使用大模型评估问题 '{question}' 的 {len(chunks)} 个知识块")
    
    results = []
    
    for i, chunk in enumerate(chunks):
        print(f"评估第 {i+1}/{len(chunks)} 个知识块...")
        
        # 添加重试逻辑
        max_retries = 3
        retry_count = 0
        while retry_count < max_retries:
            try:
                score, reason = evaluate_chunk_with_llm(question, chunk)
                results.append((chunk, score, reason))
                print(f"知识块评分: {score}, 理由: {reason[:30]}...")
                break
            except Exception as e:
                retry_count += 1
                if retry_count < max_retries:
                    wait_time = retry_count * 2  # 指数退避
                    print(f"评分失败，{wait_time}秒后重试 ({retry_count}/{max_retries})...")
                    time.sleep(wait_time)
                else:
                    print(f"评分失败，已达到最大重试次数: {str(e)}")
                    results.append((chunk, 0.0, f"评分失败: {str(e)}"))
    
    print(f"完成对 {len(results)} 个知识块的评分")
    return results

def call_qwen_rerank_api(question: str, passages: List[str], verify_ssl: bool = False, max_retries: int = 3) -> List[float]:
    """
    调用Qwen3重排模型API
    
    参数:
        question: 问题文本
        passages: 知识块列表
        verify_ssl: 是否验证SSL证书，默认为False（内网测试环境可能需要禁用）
        max_retries: 最大重试次数
        
    返回:
        重排分数列表
    """
    # API接口地址
    url = "https://aimpapi.midea.com/t-aigc/d-edp-qwen3-reranker-4b-torch/v2/models/qwen3-reranker-4b-torch/versions/1/infer"
    
    # 请求头
    headers = {
        'Content-Type': 'application/json',
        'Authorization': 'Bearer msk-c077a0a4d66af8582013b8dc731e96937d5625eb679ddec785d1f588fb0d9132'
    }
    
    # 准备请求体
    payload = {
        "inputs": [
            {
                "name": "query",
                "shape": [1],
                "datatype": "BYTES",
                "data": [question]
            },
            {
                "name": "candidates",
                "shape": [len(passages)],
                "datatype": "BYTES",
                "data": passages
            }
        ],
        "outputs": [
            {
                "name": "scores"
            }
        ]
    }
    
    retry_count = 0
    while retry_count < max_retries:
        try:
            # 发送请求
            print(f"使用Qwen3重排模型为问题 '{question}' 重排 {len(passages)} 个知识块 (尝试 {retry_count+1}/{max_retries})")
            
            # 创建带有重试机制的会话
            session = create_retry_session(retries=2)
            response = session.post(
                url, 
                headers=headers, 
                json=payload, 
                verify=verify_ssl,
                timeout=30  # 设置超时时间
            )
            response.raise_for_status()
            
            # 解析响应
            result = response.json()
            
            # 提取分数
            scores = result.get("outputs", [])[0].get("data", [])
            scores = [float(score) for score in scores]
            
            print(f"Qwen3重排完成，得分范围: [{min(scores) if scores else 0} - {max(scores) if scores else 0}]")
            return scores
            
        except requests.exceptions.SSLError as e:
            print(f"Qwen3重排SSL错误 (尝试 {retry_count+1}/{max_retries}): {str(e)}")
            retry_count += 1
            if retry_count < max_retries:
                # 指数退避
                wait_time = 2 ** retry_count
                print(f"等待 {wait_time} 秒后重试...")
                time.sleep(wait_time)
            else:
                raise RuntimeError(f"Qwen3重排API调用失败：已达到最大重试次数，SSL错误: {str(e)}")
                
        except Exception as e:
            print(f"Qwen3重排API调用异常 (尝试 {retry_count+1}/{max_retries}): {str(e)}")
            retry_count += 1
            if retry_count < max_retries:
                # 指数退避
                wait_time = 2 ** retry_count
                print(f"等待 {wait_time} 秒后重试...")
                time.sleep(wait_time)
            else:
                raise RuntimeError(f"Qwen3重排API调用失败：已达到最大重试次数，最后错误: {str(e)}")
    
    # 如果所有重试都失败，抛出异常
    raise RuntimeError("Qwen3重排API调用失败：未知错误")

def call_bge_rerank_api(question: str, passages: List[str], verify_ssl: bool = False, max_retries: int = 3) -> List[float]:
    """
    调用BGE重排模型API
    
    参数:
        question: 问题文本
        passages: 知识块列表
        verify_ssl: 是否验证SSL证书，默认为False（内网测试环境可能需要禁用）
        max_retries: 最大重试次数
        
    返回:
        重排分数列表
    """
    # API接口地址
    url = "https://aimpapi.midea.com/t-aigc/d-edp-bge-reranker-torch/v2/models/bge-reranker-v2-m3-torch/versions/1/infer"
    
    # 请求头
    headers = {
        'Content-Type': 'application/json',
        'Authorization': 'Bearer msk-8508235a080cc6a1d2017dbe066fcfeb6bedf3ee0dca4a17da00d40577b828be'
    }
    
    # 准备请求体
    payload = {
        "inputs": [
            {
                "name": "query",
                "shape": [1],
                "datatype": "BYTES",
                "data": [question]
            },
            {
                "name": "candidates",
                "shape": [len(passages)],
                "datatype": "BYTES",
                "data": passages
            }
        ],
        "outputs": [
            {
                "name": "scores"
            }
        ]
    }
    
    retry_count = 0
    while retry_count < max_retries:
        try:
            # 发送请求
            print(f"使用BGE重排模型为问题 '{question}' 重排 {len(passages)} 个知识块 (尝试 {retry_count+1}/{max_retries})")
            
            # 创建带有重试机制的会话
            session = create_retry_session(retries=2)
            response = session.post(
                url, 
                headers=headers, 
                json=payload, 
                verify=verify_ssl,
                timeout=30  # 设置超时时间
            )
            response.raise_for_status()
            
            # 解析响应
            result = response.json()
            
            # 提取分数
            scores = result.get("outputs", [])[0].get("data", [])
            scores = [float(score) for score in scores]
            
            print(f"BGE重排完成，得分范围: [{min(scores) if scores else 0} - {max(scores) if scores else 0}]")
            return scores
            
        except requests.exceptions.SSLError as e:
            print(f"BGE重排SSL错误 (尝试 {retry_count+1}/{max_retries}): {str(e)}")
            retry_count += 1
            if retry_count < max_retries:
                # 指数退避
                wait_time = 2 ** retry_count
                print(f"等待 {wait_time} 秒后重试...")
                time.sleep(wait_time)
            else:
                raise RuntimeError(f"BGE重排API调用失败：已达到最大重试次数，SSL错误: {str(e)}")
                
        except Exception as e:
            print(f"BGE重排API调用异常 (尝试 {retry_count+1}/{max_retries}): {str(e)}")
            retry_count += 1
            if retry_count < max_retries:
                # 指数退避
                wait_time = 2 ** retry_count
                print(f"等待 {wait_time} 秒后重试...")
                time.sleep(wait_time)
            else:
                raise RuntimeError(f"BGE重排API调用失败：已达到最大重试次数，最后错误: {str(e)}")
    
    # 如果所有重试都失败，抛出异常
    raise RuntimeError("BGE重排API调用失败：未知错误")

def generate_mock_scores(passages: List[str], base_score: float = 0.9, decay: float = 0.02) -> List[float]:
    """
    生成模拟的重排分数
    
    参数:
        passages: 知识块列表
        base_score: 基础分数
        decay: 分数衰减率
        
    返回:
        模拟的重排分数列表
    """
    import random
    
    scores = []
    for i in range(len(passages)):
        # 添加一些随机性
        random_factor = random.uniform(-0.05, 0.05)
        score = max(0, min(1, base_score - (i * decay) + random_factor))
        scores.append(round(score, 4))
    
    # 打乱顺序，模拟真实情况
    indices = list(range(len(scores)))
    random.shuffle(indices)
    
    shuffled_scores = [scores[i] for i in indices]
    print("使用模拟分数替代API调用结果")
    return shuffled_scores

def get_raw_search_results(question: str, top_k: int = 50) -> Dict[str, Any]:
    """
    调用检索API获取原始搜索结果（包含完整的API响应）
    
    参数:
        question: 问题文本
        top_k: 返回的知识块数量
        
    返回:
        完整的API响应结果
    """
    # API接口地址
    url = "https://iop.midea.com/rag-flow/api/search/keyword"
    
    # 请求头
    headers = {
        'appId': 'iop-knowledge',
        'appChannel': 'test',
        'Content-Type': 'application/json'
    }
    
    # 准备请求体
    payload = {
        "requestId": generate_request_id(question[:8]),
        "query": question,
        "languageList": ["zh", "en"],
        "userName": "chenmy2",
        "topk": top_k,
        "rerankMethod": "none",
        "skillList": [
            {
                "recallTypeList": ["text", "vector"],
                "skillId": "SKILL_CS31SK20250423000000192220",
                "topk": 15,
                "sourceSystem": "knowledge",
                "threshold": 0.0,
                "conditionItem": None,
                "faqAnswerConditionItem": None,
                "catalogIdList": []
            },
            {
                "recallTypeList": ["text", "vector"],
                "skillId": "SKILL_CS31SK20250418000000124477",
                "topk": 15,
                "sourceSystem": "ihr_selfhelp",
                "threshold": 0.0,
                "conditionItem": None,
                "faqAnswerConditionItem": None,
                "catalogIdList": []
            },
            {
                "recallTypeList": ["text", "vector"],
                "skillId": "SKILL_CS30SK20250417000000176056",
                "topk": 15,
                "sourceSystem": "official_doc",
                "threshold": 0.0,
                "conditionItem": None,
                "faqAnswerConditionItem": None,
                "catalogIdList": []
            },
            {
                "recallTypeList": ["text", "vector"],
                "skillId": "SKILL_CS30SK20250417000000182249",
                "topk": 15,
                "sourceSystem": "g-qms",
                "threshold": 0.0,
                "conditionItem": None,
                "faqAnswerConditionItem": None,
                "catalogIdList": []
            },
            {
                "recallTypeList": ["text", "vector"],
                "skillId": "1889572155196022785",
                "topk": 15,
                "sourceSystem": "knowledge_book",
                "threshold": 0.0,
                "conditionItem": None,
                "faqAnswerConditionItem": None,
                "catalogIdList": []
            },
            {
                "recallTypeList": ["text", "vector"],
                "skillId": "1879815395340005377",
                "topk": 15,
                "sourceSystem": "knowledge_book",
                "threshold": 0.0,
                "conditionItem": None,
                "faqAnswerConditionItem": None,
                "catalogIdList": []
            }
        ],
        "knowTypeList": None,
        "unAuthFilter": True,
        "faqSameKnowNeedOne": True,
        "lexiconGeneralizationList": None
    }
    
    try:
        # 创建带有重试机制的会话
        session = create_retry_session(retries=3)
        
        # 发送请求
        print(f"为问题 '{question}' 检索完整结果")
        response = session.post(url, headers=headers, json=payload, verify=False, timeout=30)
        response.raise_for_status()  # 如果请求失败，抛出异常
        
        # 返回完整响应
        return response.json()
        
    except Exception as e:
        print(f"API调用异常: {str(e)}")
        return {"code": "-1", "msg": str(e), "data": []} 



if __name__ == "__main__":
    # 测试检索功能
    question = "大湾区退税政策是什么？"
    top_k = 10
    result = search_knowledge_chunks(question, top_k)
    print(result)
    
    # 测试评分功能
    if result and len(result) > 0:
        # 为了测试，只使用前2个知识块
        test_chunks = result[:2]
        print("\n===测试评分API===")
        print(f"使用前 {len(test_chunks)} 个知识块进行评分测试")
        scores = evaluate_chunks_batch(question, test_chunks)
        for chunk, score, reason in scores:
            print(f"评分: {score}")
            print(f"理由: {reason}")
            print(f"知识块: {chunk[:100]}...\n")
    
    # 测试Qwen3重排功能
    if result and len(result) > 0:
        print("\n===测试Qwen3重排API===")
        rerank_scores = call_qwen_rerank_api(question, result[:3], verify_ssl=False)
        for i, (chunk, score) in enumerate(zip(result[:3], rerank_scores)):
            print(f"重排分数: {score}")
            print(f"知识块: {chunk[:100]}...\n")
    
    # 测试BGE重排功能
    if result and len(result) > 0:
        print("\n===测试BGE重排API===")
        rerank_scores = call_bge_rerank_api(question, result[:3], verify_ssl=False)
        for i, (chunk, score) in enumerate(zip(result[:3], rerank_scores)):
            print(f"重排分数: {score}")
            print(f"知识块: {chunk[:100]}...\n")